"""
包含手势分割，角点检测方法的模块
"""
import cv2
import config
import CKA
import numpy as np


FAST_RADIUS = config.FAST_RADIUS
KERNEL_SIZE = config.KERNEL_SIZE
SET_MAX_LENGTH = config.SET_MAX_LENGTH
RESCALE_RATE = config.RESCALE_RATE
REVISE_WIDTH = config.REVISE_WIDTH
REVISE_HEIGHT = config.REVISE_HEIGHT
MODEL_WIDTH = config.MODEL_WIDTH
MODEL_HEIGHT = config.MODEL_HEIGHT


# 获取手势的轮廓区域
def get_skin_area(frame):
    cv2.imshow('frame', frame)
    YCrCb = cv2.cvtColor(frame, cv2.COLOR_BGR2YCrCb)
    Y, Cr, Cb = cv2.split(YCrCb)

    ret, Cr1 = cv2.threshold(Cr, 137, 255, cv2.THRESH_BINARY)       # 大于137变白，否则变黑
    ret, Cr2 = cv2.threshold(Cr, 175, 255, cv2.THRESH_TOZERO_INV)   # 大于175变黑，否则变白
    Cr = cv2.bitwise_and(Cr1, Cr2)

    ret, Cb1 = cv2.threshold(Cb, 100, 255, cv2.THRESH_BINARY)       # 大于100变白，否则变黑
    ret, Cb2 = cv2.threshold(Cb, 118, 255, cv2.THRESH_TOZERO_INV)   # 大于118变黑，否则变白
    Cb = cv2.bitwise_and(Cb1, Cb2)

    skin = cv2.bitwise_and(Cr, Cb)
    ret, skin = cv2.threshold(skin, 0, 255, cv2.THRESH_BINARY)      # 二值化
    return skin


# 获取最大轮廓
def get_max_contour(contours):
    max_area = 0
    max_index = -1
    for i in range(len(contours)):
        if cv2.contourArea(contours[i]) > max_area:
            max_area = cv2.contourArea(contours[i])
            max_index = i
    return max_index


# 获取FAST角点
def get_fast_corner(frame, radius=10):
    fast = cv2.FastFeatureDetector_create(radius)
    fast.setNonmaxSuppression(0)
    key_points = fast.detect(frame, None)
    return key_points


# 获取手的肤色区域与FAST点集
def get_corner_and_area(mat, dex=False):
    # 获取手势区域
    copy_mat = mat.copy()
    # ptr, skin = get_skin_area(mat, dex)
    skin = get_skin_area(mat)
    # release(ptr)
    # 创建全黑图
    binary = np.zeros((copy_mat.shape[0], copy_mat.shape[1]))
    # 获取轮廓
    contours, hierarchy = cv2.findContours(skin, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if contours is not None:
        index = get_max_contour(contours)    # 获取最大轮廓
        cv2.drawContours(binary, contours, index, (255, 255, 255), thickness=-1)    # 绘制最大轮廓，将轮廓内部填充为白色
        # 创建三通道图
        binary_triple = copy_mat.copy()
        binary_triple[:, :, 0] = binary
        binary_triple[:, :, 1] = binary
        binary_triple[:, :, 2] = binary
        skin_area = cv2.bitwise_and(copy_mat, binary_triple)    # 与运算提取肤色区域
    else:
        skin_area = binary
    return skin_area


# 对图像进行切片
def cut_image(image, points, dex=False):
    cut_list = []
    out = []
    for i in range(len(points)):
        # 图像切片
        x1 = int(points[i][1] * RESCALE_RATE) - int(MODEL_WIDTH)
        y1 = int(points[i][0] * RESCALE_RATE) - int(MODEL_HEIGHT)
        x2 = int(points[i][1] * RESCALE_RATE) + int(MODEL_WIDTH)
        y2 = int(points[i][0] * RESCALE_RATE) + int(MODEL_HEIGHT)
        cut = image[x1: x2, y1: y2]
        condition = cut.shape[0] == MODEL_WIDTH * 2 and cut.shape[1] == MODEL_HEIGHT * 2
        if condition:
            cut = cv2.resize(cut, (MODEL_WIDTH, MODEL_HEIGHT))
            cut = get_corner_and_area(cut, dex)
            cut_list.append(cut)
            out.append(points[i])
    return cut_list, out


# 获取热力图和聚类中心
def get_heatmap(mat):
    background = np.zeros((mat.shape[1], mat.shape[0]))
    kp = get_fast_corner(mat, FAST_RADIUS)
    kep_points = cv2.KeyPoint_convert(kp)
    points = []
    for i in range(len(kep_points)):
        points.append([int(kep_points[i][0]), int(kep_points[i][1])])
    heat_map, points = CKA.weighted_cluster(background, points, ks=KERNEL_SIZE, max_length=SET_MAX_LENGTH)
    heat_map = heat_map.transpose()
    return heat_map, points
